package com.starsoft.verifycode.ocr;

import java.awt.image.BufferedImage;
import java.io.File;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.starsoft.verifycode.feature.FeatureLib;
import com.starsoft.verifycode.feature.SimpleFeature;
import com.starsoft.verifycode.filter.GrayFilter;
import com.starsoft.verifycode.split.LineSplit;
import com.starsoft.verifycode.util.ImageIOHelper;

public class SimpleOcr {

	Logger logger = LoggerFactory.getLogger(this.getClass());
	private List<FeatureLib> libs = new ArrayList<FeatureLib>();

	public SimpleOcr() {
	}
	
	public SimpleOcr(File libFile) {
		libs = new SimpleFeature().loadLibs(libFile);
	}

	public SimpleOcr(InputStream libIs) {
		libs = new SimpleFeature().loadLibs(libIs);
	}
	
	public List<String> getFeature(File file) {
		BufferedImage image = ImageIOHelper.read(file);
		return getFeature(image);
	}

	public List<String> getFeature(BufferedImage image) {
		BufferedImage filterImage = new GrayFilter(image).median().binarization(140).getImage();
		List<BufferedImage> splitImage = new LineSplit().split(filterImage);

		List<String> feature = new ArrayList<String>();
		for (BufferedImage tempImg : splitImage) {
			feature.add(new SimpleFeature().get(tempImg));
		}
		return feature;
	}

	public String getOcr(File file) {
		if( libs==null || libs.size()==0){
			throw new RuntimeException("没有特征数据");
		}
		String result = "";
		List<String> features = getFeature(file);
		for (int i = 0; i < features.size(); i++) {
			result += new SimpleFeature().match(features.get(i), libs);
		}
		return result;
	}

	public String getOcr(BufferedImage image) {
		if( libs==null || libs.size()==0){
			throw new RuntimeException("没有特征数据");
		}
		String result = "";
		List<String> features = getFeature(image);
		for (int i = 0; i < features.size(); i++) {
			result += new SimpleFeature().match(features.get(i), libs);
		}
		return result;
	}
	
	public List<FeatureLib> feedImageDir(File dir) {
		File[] files = dir.listFiles();
		for (File file : files) {
			String name = file.getName();
			if (name.endsWith(".jpg") || name.endsWith(".png") || name.endsWith(".gif") || name.endsWith(".jpeg")) {
				List<String> feature = this.getFeature(file);
				String names = name.substring(0, name.indexOf("."));
				if (names.length() != feature.size()) {
					logger.error("Failed to train " + name + ", number not match");
					continue;
				}
				for (int i = 0; i < names.length(); i++) {
					FeatureLib lib = new FeatureLib(names.charAt(i) + "", feature.get(i));
					libs.add(lib);
				}
			}
		}
		libs = new SimpleFeature().compressLib(libs);
		return libs;
	}

	public void loadLibs(String libStr){
		libs = new SimpleFeature().loadLibs(libStr);
	}
	
	public String getLibStr(){
		return new SimpleFeature().getLibStr(libs);
	}
	
	public List<FeatureLib> getLibs() {
		return libs;
	}

	public void setLibs(List<FeatureLib> libs) {
		this.libs = libs;
	}
}
